package mr

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"strconv"
	"sync"
	"time"
)

type TaskStyle int

const (
	Map    TaskStyle = 0
	Reduce TaskStyle = 1
)
const (
	taskTimeout = 20 * time.Second
)

type TaskState int

// Task 三种状态
const (
	Idle      TaskState = 0
	InProcess TaskState = 1
	Completed TaskState = 2
)

type Task struct {
	Loc     string
	State   TaskState
	StartUP time.Time
	ID      int
	Style   TaskStyle
}

type Coordinator struct {
	// Your definitions here.
	mu      sync.Mutex
	tasks   []*Task
	cond    *sync.Cond
	nReduce int
	buckets map[string][]*Task
}

// Your code here -- RPC handlers for the worker to call.
func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
	reply.Y = args.X + 1
	return nil
}

//
// an example RPC handler.
//
// the RPC argument and reply types are defined in rpc.go.
//
// func (c *Coordinator) Example(args *ExampleArgs, reply *ExampleReply) error {
// 	reply.Y = args.X + 1
// 	return nil
// }

//
// start a thread that listens for RPCs from worker.go
//
func (c *Coordinator) server() {
	rpc.Register(c)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := coordinatorSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

func toString(style TaskStyle, state TaskState) string {
	return strconv.Itoa(int(style)) + "-" + strconv.Itoa(int(state))
}

//
// main/mrcoordinator.go calls Done() periodically to find out
// if the entire job has finished.
//
func (c *Coordinator) Done() bool {

	// Your code here.
	c.mu.Lock()
	defer c.mu.Unlock()
	return len(c.buckets[toString(Reduce, Idle)])+len(c.buckets[toString(Reduce, InProcess)]) == 0

}

func (c *Coordinator) initTask(files []string, nReduce int) {
	// 生成map任务

	for i, file := range files {
		t := &Task{
			Loc:     file,
			State:   Idle,
			StartUP: time.Now(),
			ID:      i,
			Style:   Map,
		}
		c.tasks = append(c.tasks, t)
		c.buckets[toString(Map, Idle)] = append(c.buckets[toString(Map, Idle)], t)

	}

	// 生成reduce任务
	for i := 0; i < nReduce; i++ {
		loc := ""
		for j := 0; j < len(files); j++ {
			loc = loc + ";" + fmt.Sprintf("mr-%d-%d", j, i)
		}
		t := &Task{
			Loc:     loc,
			State:   Idle,
			StartUP: time.Now(),
			ID:      i,
			Style:   Reduce,
		}
		c.tasks = append(c.tasks, t)
		c.buckets[toString(Reduce, Idle)] = append(c.buckets[toString(Reduce, Idle)], t)
	}

}
func (c *Coordinator) refreshBucket() {
	c.buckets = make(map[string][]*Task)
	for _, v := range c.tasks {
		switch {
		case v.State == Idle:
			c.buckets[toString(v.Style, Idle)] = append(c.buckets[toString(v.Style, Idle)], v)
		case v.State == InProcess:
			c.buckets[toString(v.Style, InProcess)] = append(c.buckets[toString(v.Style, InProcess)], v)
		}
	}
}

func (t *Task) updateState(s TaskState) {

	t.State = s
}

func (t *Task) updateStartUP() {

	t.StartUP = time.Now()
}

func (t *Task) Copy() Task {
	return Task{
		Loc:     t.Loc,
		State:   t.State,
		StartUP: t.StartUP,
		ID:      t.ID,
		Style:   t.Style,
	}
}

func (c *Coordinator) hasFreeTask() bool {
	return len(c.buckets[toString(Reduce, Idle)])+len(c.buckets[toString(Reduce, InProcess)]) == 0
}

func (c *Coordinator) GetIdleTask(args *GetArgs, reply *GetReply) error {

	c.cond.L.Lock()
	defer c.cond.L.Unlock()
	defer c.refreshBucket()

	for c.hasFreeTask() {
		c.cond.Wait()
	}

	if len(c.buckets[toString(Map, Idle)]) > 0 {
		t := c.buckets[toString(Map, Idle)][0]
		reply.Task = t.Copy()
		reply.NReduce = c.nReduce
		t.updateStartUP()
		t.updateState(InProcess)
		return nil
	}
	if len(c.buckets[toString(Map, InProcess)]) > 0 {
		for _, t := range c.buckets[toString(Map, InProcess)] {
			if time.Since(t.StartUP) > taskTimeout {
				reply.Task = t.Copy()
				reply.NReduce = c.nReduce
				t.updateStartUP()
				t.updateState(InProcess)
				return nil
			}
		}
		time.Sleep(time.Second * 1)
	}

	if len(c.buckets[toString(Map, Idle)])+len(c.buckets[toString(Map, InProcess)]) == 0 {
		if len(c.buckets[toString(Reduce, Idle)]) > 0 {
			t := c.buckets[toString(Reduce, Idle)][0]
			reply.Task = t.Copy()
			reply.NReduce = c.nReduce
			t.updateStartUP()
			t.updateState(InProcess)
			return nil
		}
		if len(c.buckets[toString(Reduce, InProcess)]) > 0 {
			for _, t := range c.buckets[toString(Reduce, InProcess)] {
				if time.Since(t.StartUP) > taskTimeout {
					reply.Task = t.Copy()
					reply.NReduce = c.nReduce
					t.updateStartUP()
					t.updateState(InProcess)
					return nil
				}
			}
			time.Sleep(time.Second * 1)
		}

	}

	return nil

}
func (c *Coordinator) PutTask(args *PutArgs, reply *PutReply) error {
	c.mu.Lock()
	defer c.mu.Unlock()
	defer c.refreshBucket()
	if args.Task.Loc == "" {
		return nil
	}

	for _, v := range c.tasks {
		if v.ID == args.Task.ID && v.Style == args.Task.Style {
			if v.State == Completed {
				return nil
			}
			v.updateState(Completed)
		}
	}
	c.cond.Broadcast()
	return nil
}

//
// create a Coordinator.
// main/mrcoordinator.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeCoordinator(files []string, nReduce int) *Coordinator {
	bucket := make(map[string][]*Task)

	c := Coordinator{
		nReduce: nReduce,
		buckets: bucket,
	}
	cond := sync.NewCond(&c.mu)
	c.cond = cond

	// Your code here.
	c.initTask(files, nReduce)

	c.server()
	return &c
}
